/* Copyright (C) 2005, 2006, 2007, 2008, 2009, 2010 Mario Orsi
   This file is part of Brahms.
   Brahms is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
   Brahms is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
   You should have received a copy of the GNU General Public License along with Brahms.  If not, see <http://www.gnu.org/licenses/>. */

/* epp.c -- Direct evaluation of the electrostatic potential profile across the bilayer.
   Profiles are described by histograms along the lines of rdf histograms.
   Memory is allocated in module startUp.c, function AllocArrays(). */

#include "dataStructs.h"

extern const VecR *siteCoords, *siteOrientVec, *atomCoordinates; 
extern const RMat *siteRotMat;
extern const Atom *atom;
extern const Site *site; 
extern const VecR region; 
extern const int nAtoms, nSites, nSolutes, stepCount, sizeHistEpp, limitEpp;
extern const real deltaT, timeNow, eppDeltaZ; // bin size along the bilayer normal (z-axis) computed in step.c
extern const real watDipole, cholCharge, phosCharge, glyDipole, estDipole; 
extern int countEpp, halfSizeHistEpp; // globals, defined & initialised in startUp.c
extern real meanEppSlabVol, meanEppSlabArea, meanEppDeltaZ, meanEppBoxHeight; // globals, defined and initialised in startUp.c

/* global variables */
real *soluteHistEpp, *watHistEpp, *headgroupHistEpp, *glycerolHistEpp, *esterHistEpp, *totalHistEpp;

/* local variables */
const static real DIPOLE_Cm = CHARGE_C * LENGTH_IN_m;

// real DIPOLE_Cm = sqrt( 4. * PI * EPSILON0_IN_F_OVER_m * Cube( LENGTH_IN_m ) * ENERGY_J );

void ZeroEppHistograms();  // defined in startUp.c

static void PrintWaterDipoleMoment() // water dipole moment along the membrane normal
{
  real distanceFromBilayerCentre;
  int bin;

  FILE *eppFilePtr;
  char eppSnapFile[ 20 ];
  char time_ns[ 9 ]; // string containing timeNow, e.g. "100", in ns

  sprintf( time_ns, "%d", ( int ) ( 0.5 + timeNow *  TIME_ns ) );
  eppSnapFile[ 0 ] = '\0';
  strcat( eppSnapFile, "watDipoleMoment-" );
  strcat( eppSnapFile, time_ns );
  strcat( eppSnapFile, ".dat" );
  eppFilePtr = fopen( eppSnapFile, "w" );

  for ( bin = 0; bin < sizeHistEpp; bin++ ) { 
    distanceFromBilayerCentre = ( bin + 0.5 ) * meanEppDeltaZ - meanEppBoxHeight / 2.; 
    fprintf( eppFilePtr, "%7.3f", ( ( distanceFromBilayerCentre ) * LENGTH_nm ) ); // printing abscissa [nm]
    fprintf( eppFilePtr, "%20.9f", ( watHistEpp[ bin ] * DIPOLE_D ) ); // dipole moment [D]
    fprintf( eppFilePtr, "\n" );
  }
  fclose( eppFilePtr );
}

static void PrintElectricField()
{
  real distanceFromBilayerCentre;
  int bin;

  FILE *eppFilePtr;
  char eppSnapFile[ 20 ];
  char time_ns[ 9 ]; // string containing timeNow, e.g. "100", in ns

  sprintf( time_ns, "%d", ( int ) ( 0.5 + timeNow * TIME_ns ) );
  eppSnapFile[ 0 ] = '\0';
  strcat( eppSnapFile, "electricField-" );
  strcat( eppSnapFile, time_ns );
  strcat( eppSnapFile, ".dat" );
  eppFilePtr = fopen( eppSnapFile, "w" );

  for ( bin = 0; bin < sizeHistEpp; bin++ ) { 
    distanceFromBilayerCentre = ( bin + 0.5 ) * meanEppDeltaZ - meanEppBoxHeight / 2.; 
    fprintf( eppFilePtr, "%7.3f", ( ( distanceFromBilayerCentre ) * LENGTH_nm ) ); // printing abscissa [nm]
    fprintf( eppFilePtr, "%30.f", totalHistEpp[ bin ] * DIPOLE_Cm / ( EPSILON0_IN_F_OVER_m * VOLUME_m3 ) ); // [ V / m]
    fprintf( eppFilePtr, "%30.f", watHistEpp[ bin ]       * DIPOLE_Cm / ( EPSILON0_IN_F_OVER_m*VOLUME_m3 ) ); // electric field [ V / m]
    fprintf( eppFilePtr, "%30.f", headgroupHistEpp[ bin ] * DIPOLE_Cm / ( EPSILON0_IN_F_OVER_m*VOLUME_m3 ) ); // electric field [ V / m]
    fprintf( eppFilePtr, "%30.f", glycerolHistEpp[ bin ]  * DIPOLE_Cm / ( EPSILON0_IN_F_OVER_m*VOLUME_m3 ) ); // electric field [ V / m]
    fprintf( eppFilePtr, "%30.f", esterHistEpp[ bin ]  * DIPOLE_Cm / ( EPSILON0_IN_F_OVER_m*VOLUME_m3 ) ); // electric field [ V / m]
    if (nSolutes) fprintf( eppFilePtr, "%30.f", soluteHistEpp[bin]*DIPOLE_Cm/EPSILON0_IN_F_OVER_m*VOLUME_m3); // electric field [ V / m]
    fprintf( eppFilePtr, "\n" );
  }
  fclose( eppFilePtr );
}

static void PrintEpp() // print to file the electrical potential profile (epp)
{
  real distanceFromBilayerCentre;
  int bin;

  FILE *eppFilePtr;
  char eppSnapFile[ 20 ];
  char time_ns[ 9 ]; // string containing timeNow, e.g. "100", in ns

  sprintf( time_ns, "%d", ( int ) ( 0.5 + timeNow *  TIME_ns ) );
  eppSnapFile[ 0 ] = '\0';
  strcat( eppSnapFile, "epp-" );
  strcat( eppSnapFile, time_ns );
  strcat( eppSnapFile, ".dat" );
  eppFilePtr = fopen( eppSnapFile, "w" );

  for ( bin = 0; bin < sizeHistEpp; bin++ ) { 
    distanceFromBilayerCentre = ( bin + 0.5 ) * meanEppDeltaZ - meanEppBoxHeight / 2.; 
    fprintf( eppFilePtr, "%7.3f", ( ( distanceFromBilayerCentre ) * LENGTH_nm ) ); // printing abscissa [nm]
    fprintf( eppFilePtr, "%20.9f", totalHistEpp[ bin ]     * DIPOLE_Cm * LENGTH_IN_m / (EPSILON0_IN_F_OVER_m*VOLUME_m3) ); // [ V ]
    fprintf( eppFilePtr, "%20.9f", watHistEpp[ bin ]       * DIPOLE_Cm * LENGTH_IN_m / (EPSILON0_IN_F_OVER_m*VOLUME_m3) ); // [ V ]
    fprintf( eppFilePtr, "%20.9f", headgroupHistEpp[ bin ] * DIPOLE_Cm * LENGTH_IN_m / (EPSILON0_IN_F_OVER_m*VOLUME_m3) ); // [ V ]
    fprintf( eppFilePtr, "%20.9f", glycerolHistEpp[ bin ]  * DIPOLE_Cm * LENGTH_IN_m / (EPSILON0_IN_F_OVER_m*VOLUME_m3) ); // [ V ]
    fprintf( eppFilePtr, "%20.9f", esterHistEpp[ bin ]     * DIPOLE_Cm * LENGTH_IN_m / (EPSILON0_IN_F_OVER_m*VOLUME_m3) ); // [ V ]
    if ( nSolutes ) fprintf( eppFilePtr, "%20.9f", soluteHistEpp[bin]*DIPOLE_Cm*LENGTH_IN_m/(EPSILON0_IN_F_OVER_m*VOLUME_m3)); // [ V ]
    fprintf( eppFilePtr, "\n" );
  }
  fclose( eppFilePtr );
}

static void IntegrateHistogram( real *histArray, const real delta, const int size ) 
{
  int bin;
  real *area;

  AllocMem( area, size, real );

  area[ 0 ] = histArray[ 0 ] * delta; // working out integral (area) of first element
  
  for ( bin = 1; bin < size; bin++ ) { // integrating the rest
    area[ bin ] = area[ bin - 1 ] + histArray[ bin ] * delta;
  }
  for ( bin = 0; bin < size; bin++ ) { // overwriting
    histArray[ bin ] = area[ bin ];
  }
}

static void AccumElectroDensities() // sum up and sort the charge and dipolar densities in parallel slabs along z
{
  real zDistAtom,
    zDistCom; // site COM z-distance
  VecR rAtomLabFrame;
  int a, n, bin; // "bin" is the histogram bin number

  DO_SITE { // accumulating charge and dipolar densities 
    zDistCom = siteCoords[ n ].z; 
    bin = zDistCom / eppDeltaZ + halfSizeHistEpp;
    switch ( site[ n ].type ) {
    case WATER_TYPE:  // accumulate dipolar density
      watHistEpp[ bin ] += watDipole * siteOrientVec[ n ].z; /* multiply by z-component of site orientation vector */ 
      break;
    case CHOLINE_TYPE: // accumulate charge density
    case AMINE_TYPE: // accumulate charge density
      headgroupHistEpp[ bin ] += cholCharge; //printf("*** cholCharge = %f\n",cholCharge);
      break;
    case PHOSPHATE_TYPE: // accumulate charge density
      headgroupHistEpp[ bin ] += phosCharge; 
      break;
    case GLYCEROL_TYPE: // accumulate dipolar density  
      glycerolHistEpp[ bin ] += glyDipole * siteOrientVec[ n ].z; /* multiply by z-component of site orientation vector */ 
      break;
    case ESTER_TYPE: // accumulate dipolar density  
      esterHistEpp[ bin ] += estDipole * siteOrientVec[ n ].z ; /* multiply by z-component of site orientation vector */ 
      break;
    case TAIL_TYPE: 
      break; // no charges here!
    case SOLUTE: // atomistic solute
      DO_ATOM { // rotate the body-frame coords (atomCoordinates[i]) to the lab-frame coords (rAtomLabFrame)	
	MVMulT( rAtomLabFrame, siteRotMat[ n ].u, atomCoordinates[ a ] ); 
	zDistAtom = zDistCom + rAtomLabFrame.z;
	bin = zDistAtom / eppDeltaZ + halfSizeHistEpp;
	soluteHistEpp[ bin ] += atom[ a ].charge; //printf("*** atom[ a ].charge = %f\n", atom[ a ].charge);
      }
      break;
    default: printf("WARNING: default reached in EvalEpp switching.\n"); 
    }
  }
}

void EvalEpp()
{
  
  int bin; // "bin" is the histogram bin number

  AccumElectroDensities(); 

  meanEppSlabVol += region.x * region.y * eppDeltaZ; /* accumulating - strictly necessary only at constant pressure to account for volume changes */
  meanEppSlabArea += region.x * region.y; 
  meanEppDeltaZ += eppDeltaZ;  // eppDeltaZ computed in SingleStep()
  meanEppBoxHeight += region.z; 
    
  ++countEpp;

  if ( limitEpp == countEpp ) {  // if happy with the amount of data collected
    for ( bin = 0; bin < sizeHistEpp; bin++ ) { // normalising to have an averaged value
      soluteHistEpp[ bin ]    /= meanEppSlabVol;
      watHistEpp[ bin ]       /= meanEppSlabVol;
      headgroupHistEpp[ bin ] /= meanEppSlabVol;
      glycerolHistEpp[ bin ]  /= meanEppSlabVol;
      esterHistEpp[ bin ]  /= meanEppSlabVol;
    }
    
    meanEppDeltaZ /= countEpp;
    meanEppBoxHeight /= countEpp;

    //    PrintHeadgroupChargeDensity();
    PrintWaterDipoleMoment(); // water dipole moment along the membrane normal

    IntegrateHistogram( headgroupHistEpp, meanEppDeltaZ, sizeHistEpp );   // integrating charge densities (headgroups!) to get the electric field
    IntegrateHistogram( soluteHistEpp, meanEppDeltaZ, sizeHistEpp );   // integrating charge densities (solutes!) to get the electric field

    // change sign to dipolar densities, as elecField = - dipolarDensity [Sokhan & Tildesley, Mol Phys, 92, 625 (1997)]
    for ( bin = 0; bin < sizeHistEpp; bin++ ) {
      watHistEpp[ bin ]      *= - 1.;
      glycerolHistEpp[ bin ] *= - 1.;
      esterHistEpp[ bin ]    *= - 1.;
    }
   
    PrintElectricField();

    // integrating the electric fields to get the potentials
    IntegrateHistogram( watHistEpp,       meanEppDeltaZ, sizeHistEpp );
    IntegrateHistogram( headgroupHistEpp, meanEppDeltaZ, sizeHistEpp );
    IntegrateHistogram( glycerolHistEpp,  meanEppDeltaZ, sizeHistEpp );
    IntegrateHistogram( esterHistEpp,     meanEppDeltaZ, sizeHistEpp );
    IntegrateHistogram( soluteHistEpp,    meanEppDeltaZ, sizeHistEpp );
  
    // change sign, because by definition: potential = - integral( electric field )
    for ( bin = 0; bin < sizeHistEpp; bin++ ) {
      watHistEpp[ bin ]       *= - 1.;
      headgroupHistEpp[ bin ] *= - 1.;
      glycerolHistEpp[ bin ]  *= - 1.;
      esterHistEpp[ bin ]     *= - 1.;
      soluteHistEpp[ bin ]    *= - 1.;
    }

    // summing up the total histogram
    for ( bin = 0; bin < sizeHistEpp; bin++ ) {
      totalHistEpp[ bin ] = watHistEpp[ bin ] + headgroupHistEpp[ bin ] + glycerolHistEpp[ bin ] + esterHistEpp[ bin ];
      if ( nSolutes )
	totalHistEpp[ bin ] += soluteHistEpp[ bin ];
    }
    
    PrintEpp();

    countEpp = 0;
    meanEppSlabVol = 0.;
    meanEppSlabArea = 0.;
    ZeroEppHistograms();
  }
}
